# @Time : 2021/8/16 16:40
# @Author : xiashuobad
# @File : metrics.py
# @Software: PyCharm
"""
/**
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
import numpy as np


class Evaluator(object):

    def __init__(self, num_class, decimals=4):
        self.num_class = num_class
        self.decimals = decimals
        self.confusion_matrix = np.zeros((self.num_class,) * 2)
        self.metrics_func: dict = {
            'acc': self.Acc,
            'pre': self.Precision,
            'mpre': self.mPrecision,
            'rec': self.Recall,
            'mrec': self.mRecall,
            'f1': self.F1_score,
            'mf1': self.mF1_score,
            'iou': self.Iou,
            'miou': self.mIou,
            'fwiou': self.Fwiou,
            'kappa': self.Kappa
        }

    def Acc(self):
        # return all class overall pixel accuracy
        # acc = (TP + TN) / (TP + TN + FP + FN)
        acc = np.diag(self.confusion_matrix).sum() / self.confusion_matrix.sum()
        return acc

    def Precision(self):
        # return each category pixel accuracy(A more accurate way to call it precision)
        # precision = TP / (TP + FP)
        pre = np.diag(self.confusion_matrix) / self.confusion_matrix.sum(axis=0)
        return pre

    def mPrecision(self):
        return np.nanmean(self.Precision())

    def Recall(self):
        # return each category pixel accuracy(A more accurate way to call it precision)
        # recall = TP / (TP + FN)
        classRecall = np.diag(self.confusion_matrix) / self.confusion_matrix.sum(axis=1)
        return classRecall

    def mRecall(self):
        return np.nanmean(self.Recall())

    def F1_score(self):
        # f1_score = 2 * precision * recall / (precision + recall)
        precision = self.Precision()
        recall = self.Recall()
        f1_score = 2 * precision * recall / (precision + recall)
        return f1_score

    def mF1_score(self):
        return np.nanmean(self.F1_score())

    def Iou(self):
        # miou = TP / ((TP + FN) + (TP + FP) - TP)
        intersection = np.diag(self.confusion_matrix)
        union = np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) - intersection
        IoU = intersection / union
        return IoU

    def mIou(self):
        return np.nanmean(self.Iou())

    def Fwiou(self):
        freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix)
        miou = np.diag(self.confusion_matrix) / (
                np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
                np.diag(self.confusion_matrix))
        FWIoU = (freq[freq > 0] * miou[freq > 0]).sum()
        return FWIoU

    def Kappa(self):
        """
        kappa指数：(po - pe) / (1 - pe)
        po 和 acc相同
        pe = 所有类别分别对应的“实际与预测数量的乘积”之总和 / 样本总数的平方
        :return:
        """
        po = self.Acc()
        pe = np.matmul(self.confusion_matrix.sum(1), self.confusion_matrix.sum(0).T) / self.confusion_matrix.sum() ** 2
        kappa = (po - pe) / (1 - pe)
        return kappa

    def cal_metrics(self, metrics: (list, tuple, str) = None):
        res = {}
        if metrics is None:
            metrics = self.metrics_func.keys()
        else:
            if isinstance(metrics, str):
                metrics = [metrics]
        for m in metrics:
            res[m] = self.value_around(self.metrics_func[m]())
        return res

    def _gen_confusion_matrix(self, imgPredict, imgLabel):
        # remove classes from unlabeled pixels in gt image and predict
        mask = (imgLabel >= 0) & (imgLabel < self.num_class)
        label = self.num_class * imgLabel[mask] + imgPredict[mask]
        count = np.bincount(label, minlength=self.num_class ** 2)
        confusionMatrix = count.reshape(self.num_class, self.num_class)
        return confusionMatrix

    def add_batch(self, imgPredict, imgLabel):
        assert imgPredict.shape == imgLabel.shape
        self.confusion_matrix += self._gen_confusion_matrix(imgPredict, imgLabel)

    def reset(self):
        self.confusion_matrix = np.zeros((self.num_class, self.num_class))

    def value_around(self, val):
        '''
        数值精度
        :param val:
        :return:
        '''
        if type(val) is np.ndarray:
            return np.around(val, decimals=self.decimals).tolist()
        return round(val, self.decimals)


if __name__ == '__main__':
    pass
